
// struct Particle {
//   pos : vec2f,
//   vel : vec2f,
// }
// @group(0) @binding(0) var<storage, read_write> particles: array<Particle>;

struct UniformParams {
  img_size: vec2<u32>,
  uv_offset: vec2<i32>,
};
@group(0) @binding(0) var<uniform> params: UniformParams;

@group(0) @binding(1) var from_tex: texture_2d<f32>;
// @group(0) @binding(1) var from_tex: texture_storage_2d<rgba8unorm, read>;
@group(0) @binding(2) var to_tex: texture_storage_2d<rgba8unorm, write>;

@compute @workgroup_size(16, 16)
fn cs_main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let uv = global_id.xy;
    let uv_i = vec2<i32>(uv);
    if (uv.x >= params.img_size.x || uv.y >= params.img_size.y) {
        return;
    }

    // let particle = particles[uv.x * uv.y];
    let offset = params.uv_offset;
    let size = params.img_size;
    let texel = textureLoad(from_tex, uv, 0);

    var r: vec4<f32> = vec4(0.0);
    var cnt: u32 = 0u;

    let lt = vec2<i32>(uv_i.x - offset.x, uv_i.y - offset.y);
    if (is_legal_uv(lt, size)) {
      r += textureLoad(from_tex, lt, 0);
      cnt++;
    }
    let lc = vec2<i32>(uv_i.x - offset.x, uv_i.y);
    if (is_legal_uv(lc, size)) {
      r += textureLoad(from_tex, lc, 0);
      cnt++;
    }
    let lb = vec2<i32>(uv_i.x - offset.x, uv_i.y + offset.y);
    if (is_legal_uv(lb, size)) {
      r += textureLoad(from_tex, lb, 0);
      cnt++;
    }
    let cb = vec2<i32>(uv_i.x, uv_i.y + offset.y);
    if (is_legal_uv(cb, size)) {
      r += textureLoad(from_tex, cb, 0);
      cnt++;
    }
    let rb = vec2<i32>(uv_i.x + offset.x, uv_i.y + offset.y);
    if (is_legal_uv(rb, size)) {
      r += textureLoad(from_tex, rb, 0);
      cnt++;
    }
    let rc = vec2<i32>(uv_i.x + offset.x, uv_i.y);
    if (is_legal_uv(rc, size)) {
      r += textureLoad(from_tex, rc, 0);
      cnt++;
    }
    let rt = vec2<i32>(uv_i.x + offset.x, uv_i.y - offset.y);
    if (is_legal_uv(rt, size)) {
      r += textureLoad(from_tex, rt, 0);
      cnt++;
    }
    let ct = vec2<i32>(uv_i.x, uv_i.y - offset.y);
    if (is_legal_uv(ct, size)) {
      r += textureLoad(from_tex, ct, 0);
      cnt++;
    }
    
    let self_w = 1.0 / 9.0;
    var rst: vec4<f32>;
    if (cnt == 0u) {
      rst = texel;
      rst.r = 0.9;
    } else if (cnt > 0u && cnt <= 5u) {
      rst = r / f32(cnt);
      rst.g = 0.9;
    } else {
      rst = r / f32(cnt);
      rst.b = 0.9;
    }
    textureStore(to_tex, uv_i, rst);
}

fn is_legal_uv(uv: vec2<i32>, size: vec2<u32>) -> bool {
  return uv.x >= 0 && uv.y >= 0 && uv.x < i32(size.x) && uv.y < i32(size.y);
} 